package com.boulderai.mcp.service;

import com.boulderai.mcp.client.McpClient;
import com.boulderai.mcp.client.McpClientFactory;
import com.boulderai.mcp.model.Tool;
import com.boulderai.mcp.model.Resource;
import com.fasterxml.jackson.databind.JsonNode;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * MCP服务层
 * 封装MCP客户端操作逻辑，支持多服务器连接管理
 */
@Service
public class McpService {

    // 管理多个MCP服务器连接
    private final Map<String, McpClient> mcpClients = new ConcurrentHashMap<>();
    // 会话到服务器URI的映射
    private final Map<String, String> sessionToServerMap = new ConcurrentHashMap<>();

    /**
     * 连接到MCP服务器
     */
    public CompletableFuture<Void> connect(String sessionId, String serverUri) {
        McpClient mcpClient = McpClientFactory.createWebSocketClient();
        mcpClients.put(serverUri, mcpClient);
        sessionToServerMap.put(sessionId, serverUri);
        return mcpClient.connect(serverUri);
    }

    /**
     * 断开指定会话的连接
     */
    public CompletableFuture<Void> disconnect(String sessionId) {
        String serverUri = sessionToServerMap.get(sessionId);
        if (serverUri != null) {
            McpClient mcpClient = mcpClients.get(serverUri);
            if (mcpClient != null) {
                sessionToServerMap.remove(sessionId);
                // 如果没有其他会话使用此服务器，则断开连接
                boolean hasOtherSessions = sessionToServerMap.containsValue(serverUri);
                if (!hasOtherSessions) {
                    mcpClients.remove(serverUri);
                    return mcpClient.disconnect();
                }
            }
        }
        return CompletableFuture.completedFuture(null);
    }

    /**
     * 检查指定会话的连接状态
     */
    public boolean isConnected(String sessionId) {
        String serverUri = sessionToServerMap.get(sessionId);
        if (serverUri != null) {
            McpClient mcpClient = mcpClients.get(serverUri);
            return mcpClient != null && mcpClient.isConnected();
        }
        return false;
    }

    /**
     * 初始化会话
     */
    public CompletableFuture<JsonNode> initialize(String sessionId, JsonNode clientInfo) {
        McpClient mcpClient = getClientForSession(sessionId);
        if (mcpClient == null) {
            return CompletableFuture.supplyAsync(() -> {
                throw new IllegalStateException("Not connected to server");
            });
        }
        return mcpClient.initialize(clientInfo);
    }

    /**
     * 获取工具列表
     */
    public CompletableFuture<List<Tool>> listTools(String sessionId) {
        McpClient mcpClient = getClientForSession(sessionId);
        if (mcpClient == null) {
            return CompletableFuture.supplyAsync(() -> {
                throw new IllegalStateException("Not connected to server");
            });
        }
        return mcpClient.listTools();
    }

    /**
     * 调用工具
     */
    public CompletableFuture<JsonNode> callTool(String sessionId, String toolName, JsonNode arguments) {
        McpClient mcpClient = getClientForSession(sessionId);
        if (mcpClient == null) {
            return CompletableFuture.supplyAsync(() -> {
                throw new IllegalStateException("Not connected to server");
            });
        }
        return mcpClient.callTool(toolName, arguments);
    }

    /**
     * 获取资源列表
     */
    public CompletableFuture<List<Resource>> listResources(String sessionId) {
        McpClient mcpClient = getClientForSession(sessionId);
        if (mcpClient == null) {
            return CompletableFuture.supplyAsync(() -> {
                throw new IllegalStateException("Not connected to server");
            });
        }
        return mcpClient.listResources();
    }

    /**
     * 读取资源
     */
    public CompletableFuture<JsonNode> readResource(String sessionId, String resourceUri) {
        McpClient mcpClient = getClientForSession(sessionId);
        if (mcpClient == null) {
            return CompletableFuture.supplyAsync(() -> {
                throw new IllegalStateException("Not connected to server");
            });
        }
        return mcpClient.readResource(resourceUri);
    }

    /**
     * 发送自定义请求
     */
    public CompletableFuture<JsonNode> sendRequest(String sessionId, String method, JsonNode params) {
        McpClient mcpClient = getClientForSession(sessionId);
        if (mcpClient == null) {
            return CompletableFuture.supplyAsync(() -> {
                throw new IllegalStateException("Not connected to server");
            });
        }
        return mcpClient.sendRequest(method, params);
    }

    /**
     * 发送通知
     */
    public void sendNotification(String sessionId, String method, JsonNode params) {
        McpClient mcpClient = getClientForSession(sessionId);
        if (mcpClient != null) {
            mcpClient.sendNotification(method, params);
        }
    }

    /**
     * 获取指定会话对应的MCP客户端
     */
    private McpClient getClientForSession(String sessionId) {
        String serverUri = sessionToServerMap.get(sessionId);
        if (serverUri != null) {
            return mcpClients.get(serverUri);
        }
        return null;
    }
}